import os
import sys
from PyQt5.QtWidgets import *
from PyQt5.QtCore import *
from PyQt5.QtGui import *
import requests
import sounddevice as sd
import wave
import struct
import numpy as np
from collections import deque
from io import BytesIO
import onnxruntime as ort
import time
import subprocess
import multiprocessing

script_dir = os.path.dirname(__file__)
#ASR_URL = 'http://oy.nmgoyun.com/api/asr_form/v1'
ASR_URL = 'http://202.207.12.156/asr'
MODEL = os.path.join(script_dir, 'vad_model.onnx')
LICENSE= '/opt/Oyun/license/license.txt'


def worker():
    subprocess.run(["/usr/share/fcitx/Oyun/OyunUI/oyun/bin/python","/usr/share/fcitx/Oyun/OyunUI/closetray.py"])

def asr1(f, sign, nonce):
    data = {'pid': 'linux',
            'sign': sign,
            'timestamp': int(time.time() * 1000),
            'nonce': nonce,
            'format': 1,
            'audio': f}
    try:
        return requests.post(ASR_URL, files=data).json()['data']['mn']
    except Exception as e:
        return str(e)
def asr(f, sign, nonce):
    data = {'wave': f}
    try:
        return requests.post(ASR_URL, files=data).json()['mongol']
    except Exception as e:
        return str(e)

class VADModel:
    def __init__(self,
                 model=MODEL,
                 trig_sum: float = 0.26,
                 neg_trig_sum: float = 0.07,
                 num_steps: int = 8,
                 num_samples_per_window: int = 4000):
        self.num_samples = num_samples_per_window
        self.num_steps = num_steps
        assert self.num_samples % num_steps == 0
        self.step = int(self.num_samples / num_steps)  # 500 samples is good enough
        self.prev = np.zeros(self.num_samples)
        self.last = False
        self.triggered = False
        self.buffer = deque(maxlen=num_steps)
        self.trig_sum = trig_sum
        self.neg_trig_sum = neg_trig_sum
        self.vad_model = ort.InferenceSession(model)
        self.model = lambda x: self.vad_model.run(None, {'input': x})[0]
        self.in_speech = False
        self.speech = []
        self.last_chunk = []

    def refresh(self):
        self.prev = np.zeros(self.num_samples)
        self.last = False
        self.triggered = False
        self.buffer = deque(maxlen=self.num_steps)
        self.in_speech = False
        self.speech = []

    def prepare_batch(self, wav_chunk):
        if len(wav_chunk) < self.num_samples:
            wav_chunk = np.concatenate([np.zeros(self.num_samples - len(wav_chunk)), wav_chunk])
            self.last = True
        stacked = np.concatenate([self.prev, wav_chunk])
        self.prev = wav_chunk
        overlap_chunks = [[stacked[i:i + self.num_samples]]
                          for i in range(self.step, self.num_samples + 1, self.step)]
        return np.concatenate(overlap_chunks)

    def state(self, model_out):
        current_speech = ''
        speech_probs = model_out[:, 1]  # this is very misleading
        for i, predict in enumerate(speech_probs):
            self.buffer.append(predict)
            if ((sum(self.buffer) / len(self.buffer)) >= self.trig_sum) and not self.triggered:
                self.triggered = True
                current_speech = 'start'
            if ((sum(self.buffer) / len(self.buffer)) < self.neg_trig_sum) and self.triggered:
                current_speech = 'end'
                self.triggered = False
        return current_speech

    def __call__(self, chunk):
        chunk = np.array(struct.unpack('h' * (len(chunk) // 2), chunk)) / 2 ** 15
        batch = self.prepare_batch(chunk)
        outs = self.model(np.array(batch, np.float32))
        state = self.state(outs)
        if state == 'start':
            self.in_speech = True
            self.speech = []
            self.speech.append(self.last_chunk)
        elif state == 'end':
            self.in_speech = False
            if len(self.speech):
                return np.concatenate(self.speech)
        if self.in_speech:
            self.speech.append(chunk)
        else:
            self.last_chunk = chunk


WAVE_FORMAT_PCM = 0x0001


def pcm2wave(pcmdata):
    nframes = len(pcmdata)
    nchannels = 1
    sampwidth = 2
    framerate = 16000

    wavfile = BytesIO()
    # write header
    wavfile.write(b'RIFF')
    wavfile.write(struct.pack('<L4s4sLHHLLHH4s',
                              36 + len(pcmdata) * 2, b'WAVE', b'fmt ', 16,
                              WAVE_FORMAT_PCM, nchannels, framerate,
                              nchannels * framerate * sampwidth,
                              nchannels * sampwidth, sampwidth * 8, b'data'))
    # write data
    wavfile.write((pcmdata * (2 ** 15)).astype('h').tobytes())
    return wavfile.getvalue()


vad = VADModel()
default_input_device = sd.default.device[0]
device_info = sd.query_devices(default_input_device)
sample_rate = device_info['default_samplerate']


class WorkerThread(QThread):
    def __init__(self, sign, nonce):
        super().__init__()
        self.sign = sign[:-1]
        self.nonce = nonce[:-1]

    updateButton = pyqtSignal(str)

    def run(self):
        self.endf = 0

        def callback(indata, frames, time, status):
            indata = indata.tobytes()
            # print(indata)
            state = vad(indata)
            print(state)
            if state is not None:
                result = asr(pcm2wave(state), self.sign, self.nonce)
                try:
                    response = requests.get('http://127.0.0.1:' + '56181' + '/symbols?value=' + result)
                except:
                    pass
                print(result)
                self.endf = 1

        # 打开音频流
        with sd.InputStream(channels=1, samplerate=16000, blocksize=8000, dtype='int16', callback=callback):
            # 持续录制音频，直到手动停止
            while True:
                if self.endf:
                    break
                try:
                    sd.sleep(100)
                except KeyboardInterrupt:
                    break


class ASRWindow(QWidget):
    # 定义常量
    WINDOW_WIDTH = 300
    WINDOW_HEIGHT = 180
    BORDER_COLOR = QColor("#02A7F0")
    BORDER_WIDTH = 1

    def __init__(self):
        super().__init__()
        # 读取license.txt文件，如果文件不存在或读取失败，则会弹出信息框提示无权限，并退出程序
        self.license = self.read_license()
        if not self.license:
            QMessageBox.information(self, '信息', '无权限')
            exit()

        # 设置窗口大小、样式和页面背景颜色
        self.resize(self.WINDOW_WIDTH, self.WINDOW_HEIGHT)
        self.setWindowFlags(Qt.WindowType.Popup | Qt.WindowType.Tool |
                            Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
        self.setStyleSheet("background-color: white;")

        # 创建开始按钮和关闭按钮，并连接点击事件
        self.create_buttons()

        # 创建一个WorkerThread类的实例对象
        self.create_thread()
        # 将窗口移动到屏幕的右下角
        screen = QDesktopWidget().screenGeometry()
        self.move(screen.width() - self.width() - 100, screen.height() - self.height() - 200)

        # 设置窗口标题
        self.setWindowTitle("   ")

        # 添加拖动窗口功能
        self.dragPosition = None
        self.is_recording = False

    def closeEvent(self, event):
        event.ignore()
        self.hide()
        #subprocess.run(["/usr/share/fcitx/Oyun/OyunUI/oyun/bin/python","/usr/share/fcitx/Oyun/OyunUI/closetray.py"])
        # p = multiprocessing.Process(target=worker)
        # p.start()

    # def showEvent(self, event):
    #     tray_available = subprocess.run(["pgrep","-f","closetray.py"], stdout=subprocess.PIPE).stdout.decode('utf-8').splitlines()
    #     for tray_id in tray_available:
    #         subprocess.run(["kill",tray_id])

    def read_license(self):
        # 读取license.txt文件，如果文件不存在或读取失败，则返回None
        try:
            with open(LICENSE, 'r') as f:
                return f.readlines()
        except:
            return None

    def create_buttons(self):
        # 创建开始按钮和关闭按钮，并连接点击事件
        self.button = QPushButton(self)
        self.button.clicked.connect(self.button_clicked)
        self.button.setGeometry(126, 100, 48, 48)
        self.button.setIconSize(QSize(48, 48))
        icon = QIcon(QPixmap(os.path.join(script_dir, f"static/images/asrImages/话筒.png")))

        # 添加文本
        self.start_label = QLabel(self)
        self.start_label.setGeometry(50, 20, 200, 60)
        pixmap = QPixmap(os.path.join(script_dir, f"static/images/asrImages/prompt.png"))
        self.start_label.setPixmap(pixmap)
        self.start_label.setScaledContents(True)
        # 添加文本
        self.stop_label = QLabel(self)
        self.stop_label.setGeometry(80, 20, 140, 60)
        pixmap = QPixmap(os.path.join(script_dir, f"static/images/asrImages/voice.png"))
        self.stop_label.setPixmap(pixmap)
        self.stop_label.setScaledContents(True)
        self.stop_label.hide()

        self.button.setIcon(icon)
        cover_color = QColor("#02A7F0").lighter(200)
        self.button.setStyleSheet('''
                    QPushButton {
                        color: transparent;
                        border: 1px solid #02A7F0;
                        border-radius: 24px;
                    }
                    QPushButton:hover {
                        background-color: %s;
                    }
                ''' % (cover_color.name()))
        # 添加gif图片
        self.label = QLabel(self)
        self.label.setGeometry(10, 98, 280, 50)
        self.movie = QMovie(os.path.join(script_dir, f"static/images/asrImages/GIF.gif"))
        self.label.setMovie(self.movie)
        self.button.raise_()
        self.movie.start()
        self.movie.setSpeed(0)

        # 关闭按钮

        self.close_button = QPushButton(self)
        self.close_button.setGeometry(self.width() - 30, 10, 20, 20)
        icon = QIcon(QPixmap(os.path.join(script_dir, f"static/images/asrImages/错误.png")))
        self.close_button.setIcon(icon)
        self.close_button.clicked.connect(self.onCloseButtonClicked)
        # 设置关闭按钮的样式表，去掉按钮的边框和阴影效果
        self.close_button.setStyleSheet('''
                    QPushButton {
                        border: 1.5px dashed #02A7F0;
                        padding: 2px;
                        background-color: transparent;
                    }
                    QPushButton:hover {
                        background-color: %s;
                    }''' % (cover_color.name()))

    def onCloseButtonClicked(self):
        self.hide()
        # subprocess.run(["/usr/share/fcitx/Oyun/OyunUI/oyun/bin/python","/usr/share/fcitx/Oyun/OyunUI/closetray.py"])

    def create_thread(self):
        # 创建一个WorkerThread类的实例对象
        self.thread = WorkerThread(self.license[0], self.license[1])

    def paintEvent(self, event):
        self.SHADOW_BLUR_RADIUS = 25
        self.radius = 10
        # 绘制边框
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)

        # 创建一个圆角矩形
        path = QPainterPath()
        path.addRoundedRect(QRectF(self.BORDER_WIDTH / 2, self.BORDER_WIDTH / 2,
                                   self.width() - self.BORDER_WIDTH, self.height() - self.BORDER_WIDTH),
                            self.radius, self.radius)


        # 绘制矩形边框
        pen = QPen(self.BORDER_COLOR)
        pen.setWidth(self.BORDER_WIDTH)
        painter.setPen(pen)
        painter.drawPath(path)

    def button_clicked(self):
        if not self.is_recording:
            # 开始录制
            self.movie.jumpToFrame(0)
            self.movie.setSpeed(100)
            self.movie.start()
            self.is_recording = True
            icon = QIcon(QPixmap(os.path.join(script_dir, f"static/images/asrImages/暂停 (1).png")))
            self.button.setIconSize(QSize(48, 48))
            cover_color = QColor("#02A7F0").lighter(200)
            self.button.setStyleSheet('''
                            QPushButton {
                                color: transparent;
                                border: 1px solid #02A7F0;
                                border-radius: 24px;
                                width: 48px;
                                height: 48px;
                            }
                            QPushButton:hover {
                                background-color: %s;
                            }
                            QPushButton::icon {
                                width: 32px;
                                height: 32px;
                                position: center; /* 图标居中显示 */
                            }
                        ''' % (cover_color.name()))
            self.button.setIcon(icon)
            self.start_label.hide()
            self.start_label.hide()
            self.stop_label.show()
            self.thread.start()

        else:
            # 停止录制
            self.movie.stop()
            self.movie.jumpToFrame(0)
            self.is_recording = False
            icon = QIcon(QPixmap(os.path.join(script_dir, f"static/images/asrImages/话筒.png")))
            self.button.setIconSize(QSize(48, 48))
            cover_color = QColor("#02A7F0").lighter(200)
            self.button.setStyleSheet('''
                            QPushButton {
                                color: transparent;
                                border: 1px solid #02A7F0;
                                border-radius: 24px;
                                width: 48px;
                                height: 48px;
                            }
                            QPushButton:hover {
                                background-color: %s;
                            }
                            QPushButton::icon {
                                width: 32px;
                                height: 32px;
                                position: center; /* 图标居中显示 */
                            }
                        ''' % (cover_color.name()))
            self.button.setIcon(icon)
            self.start_label.show()
            self.stop_label.hide()
            self.thread.endf = 1


    def mousePressEvent(self, event):
        # 处理鼠标按下事件
        if event.button() == Qt.LeftButton:
            self.dragPosition = event.globalPos() - self.frameGeometry().topLeft()
            event.accept()

    def mouseMoveEvent(self, event):
        # 处理鼠标移动事件
        if event.buttons() == Qt.LeftButton:
            self.move(event.globalPos() - self.dragPosition)
            event.accept()

    def mouseReleaseEvent(self, event):
        # 处理鼠标释放事件
        self.dragPosition = None


if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = ASRWindow()
    window.show()
    app.exec_()
